import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.nn import GATConv, Linear, to_hetero


class Classifier(nn.Module):
    def forward(self, x, index_edge_label):
        edge_feat_A = x[index_edge_label[0]]
        edge_feat_B = x[index_edge_label[1]]
        return (edge_feat_A * edge_feat_B).sum(dim=-1)


class GNN(nn.Module):
    def __init__(self, hidden_channels):
        super(GNN, self).__init__()
        self.conv1 = SAGEConv(hidden_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)
        self.conv3 = SAGEConv(hidden_channels, hidden_channels)



    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        x = self.conv3(x, edge_index)
        return x


class Model(nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.gnn = GNN(hidden_channels)
        self.classifier = Classifier()

    def change_model(self, data):
        self.gnn = to_hetero(self.gnn, metadata=data.metadata())
        # pass

    def forward(self, batch_data):
        batch_data = batch_data.to("cuda:0")
        # print("batch_databatch_databatch_databatch_databatch_databatch_databatch_data",batch_data)
        new_x = self.gnn(batch_data.x, batch_data.edge_label_index)
        out = self.classifier(new_x, batch_data.edge_label_index)
        return out